import re
import pandas as pd
import streamlit as st

from api.agent_config import create_agent
from api.model_config import create_llm_model


def get_prompt_content_var(prompt_content):
    """正则表达式模式，用于匹配大括号中的内容"""
    pattern = r"\{(.*?)\}"
    matches = re.findall(pattern, prompt_content)
    table = [{"command": key, "rating": ""} for key in matches]
    return table


def create_table(prompt_content):
    if prompt_content:
        table = get_prompt_content_var(prompt_content)
        if table:
            df = pd.DataFrame(table)
            edited_df = st.data_editor(
                df,
                column_config={
                    "command": "Key",
                    "rating": st.column_config.TextColumn(
                        "Value",
                        help="请输入目前key的对应value"
                    )
                },
                disabled=["command"],
                hide_index=True,
            )
            return edited_df
    return None


st.set_page_config(layout="wide")

# 初始化侧边栏
with st.sidebar:
    st.title("自定义GPTs对话系统")
    apply_args = st.button("应用改变(这会重启会话)")
    st.divider()

    # 对话前提示词
    prompt = st.text_area("对话前提示词：", placeholder="可自定义提示词，变量用{var}定义")
    st.write("提示词用于对AI的回复做出一系列指令和约束。可插入表单变量，如{input}。这段提示词不会被最终用户所看到")

    # 提取提示词中的变量并创建表格
    edited_df = create_table(prompt)
    st.divider()

    # 知识库挂载
    uploaded_files = st.file_uploader("挂载企业知识库", accept_multiple_files=True, type=["pdf", "txt", "db"])
    if uploaded_files:
        st.write('你选择了以下文件：')
        for uploaded_file in uploaded_files:
            st.write(uploaded_file.name)

    st.divider()

    # 工具选择
    selected_tools = st.multiselect("请选择使用工具：", ["Python解释器", "天气工具", "数据库查询工具"])

    st.divider()

    # 开场白选择
    use_opening = st.checkbox("是否需要开场白", value=True)

# 标题和模型配置按钮
st.title("自定义GPTs对话系统")
with st.popover("模型配置", help="配置OpenAI GPTs模型参数"):
    apply_model = st.button("应用改变(会重启会话)")
    api_key = st.text_input("输入密钥：", type="password")
    base_url = st.text_input("Base URL：", "https://api.aigc369.com/v1")
    model_name = st.selectbox("模型选择：", ["gpt-3.5-turbo", "gpt-4o", "gpt-4"])
    temperature = st.slider("温度选择：", 0.0, 1.0, 0.5)
    max_tokens = st.number_input("最大token选择：", min_value=1, max_value=4096, value=500)

st.divider()

if not api_key:
    st.info("请输⼊你的密钥")
    st.stop()

# 初始化模型
if apply_model or apply_args or "model" not in st.session_state:
    st.session_state["model"] = create_llm_model(api_key, base_url, model_name, temperature, max_tokens)
    print("model创建")

if apply_model or apply_args or "final_prompt" not in st.session_state:
    variable_values = {}
    if edited_df is not None:
        for index, row in edited_df.iterrows():
            variable_values[row["command"]] = row["rating"]

    st.session_state["final_prompt"] = prompt.format(**variable_values) + "\n"

if apply_model or apply_args or "agent" not in st.session_state:
    st.session_state["agent"] = create_agent(st.session_state["model"], selected_tools, uploaded_files,
                                             st.session_state["final_prompt"])

# 开场白
if apply_model or apply_args or "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "ai", "content": "你好，我是逆蝶，你有什么想说的"}]
if not use_opening:
    st.session_state["messages"] = []

for message in st.session_state["messages"]:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# 对话框
user_input = st.chat_input("我在看着你👁️")

if user_input:
    st.session_state["messages"].append({"role": "human", "content": user_input})
    with st.chat_message("human"):
        st.write(user_input)

if user_input:
    with st.spinner("AI正在思考🥵🥵🥵🥵🥵"):
        _input = "你的设定是：###设定开始" + st.session_state["final_prompt"] + "设定结束###" + user_input
        response = st.session_state["agent"].invoke({"input": _input})
        st.session_state["messages"].append({"role": "ai", "content": response["output"]})
        with st.chat_message("ai"):
            st.write(response["output"])